# coding:utf-8
from BigDataWeb.algorithm import Algorithm
import jieba
from gensim import corpora, models, similarities


class TextSimilarity(Algorithm):
    text_1 = ""
    text_2 = ""
    text_3 = ""
    similarities = []
    
    def __init__(self):
        Algorithm.__init__(self)
        self.algorithm_name = "文本相似度"
        self.ipynb_template_name = "text_similarity-template.ipynb"
    
    def implent(self):
        Algorithm.implent(self)
        doc_list = [self.text_1, self.text_2, self.text_3]
        # 分词
        texts = []
        for doc in doc_list:
            texts.append([word for word in jieba.cut(doc)])
        # 词袋
        dictionary = corpora.Dictionary(texts)
        corpus = [dictionary.doc2bow(text) for text in texts]
        # 建模
        tfidf = models.TfidfModel(corpus)
        # 相似度
        index = similarities.SparseMatrixSimilarity(tfidf[corpus], num_features=len(dictionary.keys()))
        # 整理similarities
        self.similarities = []
        for i in range(0, len(corpus)):
            self.similarities.append(index[tfidf[corpus[i]]])
    
    def prepareIpynbItems(self):
        Algorithm.prepareIpynbItems(self)
        self.ipynb_items["#text_1#"] = self.text_1
        self.ipynb_items["#text_2#"] = self.text_2
        self.ipynb_items["#text_3#"] = self.text_3
